#!/usr/bin/env python

import yaml
import socket
import struct
import itertools

def ip_to_int(ip):
    return struct.unpack('>I', socket.inet_aton(ip))[0]

def single_or_list(data):
    if type(data) == list:
        out = data
    else:
        out = [data]
    return out

class PortGroup:
    def __init__(self, ports, directionCode):
        self.ports = ports
        self.directionCode = directionCode
        self.check_for_port_negation()
        self.proto = self.get_proto()

    def check_for_port_negation(self):
        self.negatePort = False 
        if '!' in "".join(self.ports):
            self.negatePort = True

    def get_proto(self):
        proto = 'tcp'
        if '/' in "".join(self.ports):
            proto = self.ports[0].split('/')[0]

        if self.negatePort:
            proto = proto[1:]
        return proto

    def as_text(self):
        portNos = [] 
        for port in self.ports:
            if "/" in port:
                port = port.split("/")[1]
            else:
                port = port
            portNos.append(port)

        return ",".join(portNos)

    def iptables_multiport_option(self):
        out = []
        if len(self.ports) > 1:
            out = ['-m multiport']
        return out

    def iptables_proto_option(self):
        return ["-p %s" % self.proto]

    def iptables_port_negation(self):
        out = []
        if self.negatePort:
           out = ['!']
        return out

    def iptables_port_option(self):
        if len(self.ports) > 1:
            option = "--%sports " %self.directionCode
        else:
            option = "--%sport " %self.directionCode
        return [option + self.as_text()]

    def iptables_port_match_options(self):
        out = []
        if self.as_text() != "ANY":
            out = list(itertools.chain(self.iptables_multiport_option(), self.iptables_proto_option(),
                               self.iptables_port_negation(), self.iptables_port_option()))
        return out

class NetObj:
    def __init__(self, data):
        if data == "ANY":
            self.ip = 'ANY'
            self.port = 'ANY'
        else:
            self.ip = data.get('ip', 'ANY')
            self.port = data.get('port', 'ANY')

    @staticmethod
    def proto_part(port):
        if not '/' in port:
           proto = 'tcp' 
        else:
            proto = port.split('/')[0]
        return proto

    def group_ports_by_protocol(self, ports):
        return [list(portParts) for proto, portParts in itertools.groupby(ports, self.proto_part)]

    def iptables_match_options(self):
        out = []
        ports = [str(port) for port in single_or_list(self.port)]
        ips = single_or_list(self.ip)
        portGroups = self.group_ports_by_protocol(ports)

        for portGroup in portGroups:
            for ip in ips:
                out.append(self.iptables_match_options_single(ip, portGroup))
        return out

    def iptables_match_options_single(self, ip, portGroups):
        negateIP = False 
        if ip[0] == '!':
            negateIP = True
            ip = ip[1:]

        out = []
        if '-' in ip:
            out.append('-m iprange')

            if negateIP:
                out.append("!")

            out.append('--%s-range %s' %(self.extDirectionCode, ip))

        elif ip != 'ANY':
            if negateIP:
                out.append("!")

            out.append("-%s %s" %(self.directionCode, ip))

        out += PortGroup(portGroups, self.directionCode).iptables_port_match_options()

        return tuple(out)

    def __cmp__(self, other):
        return cmp(self.get_size(), other.get_size())

    def get_size(self):
        if "/" in self.ip:
            mask = int(self.ip.split("/")[1])
            size = 2**(32-mask)
        elif "-" in self.ip:
            minIP, maxIP = self.ip.split("-")
            minIPInt = ip_to_int(minIP)
            maxIPInt = ip_to_int(maxIP)
            size = abs(maxIPInt - minIPInt)
        elif self.ip == "ANY":
            size = 2**32
        else:
            size = 1
        return size

class Server(NetObj):
    directionCode = "d"
    extDirectionCode = "dst"

class Client(NetObj):
    directionCode = "s"
    extDirectionCode = "src"


class AccessLevel:
    def __init__(self, data, policy):
        self.data = data
        self.policy = policy

        serversData = data["servers"]
        self.servers = [Server(serverData) for serverData in serversData] 

        clientsData = data["clients"]
        self.clients = [Client(clientData) for clientData in clientsData] 

        self.name = data["access"]

    def chain_name(self):
        return self.name + "_servers"

    def next_servers_chain(self):
        myIdx = self.policy.levels.index(self)
        if myIdx == 0:
            return 'DROP'
        else:
            return self.policy.levels[myIdx - 1].chain_name()

    def iptables_server_chain(self):
        init = ['iptables', '-N', self.chain_name()]

        serverLines = []
        servers = self.servers[:]
        servers.sort()
        for server in servers:
            for subRule in server.iptables_match_options():
                matchArgs = []
                for matchOption in subRule: 
                    matchArgs += matchOption.split()

                serverLines.append(['iptables', '-A', self.chain_name()] + matchArgs + ['-j', 'ACCEPT'])

        cascade = ['iptables', '-A', self.chain_name(), '-j', self.next_servers_chain()]
        return [init] + serverLines + [cascade]


class Policy:

    class EmptyPolicyDocument(Exception): pass

    @classmethod
    def from_yaml(cls, yamlStr):
        data = yaml.load(yamlStr)
        return cls(data)

    def __init__(self, data):
        if not data:
            raise self.EmptyPolicyDocument

        self.levels = [AccessLevel(levelData, self) for levelData in data]

    def iptables_authorize_clients(self):
        clientLevelPairs = []
        for levelNo in range(0, len(self.levels)):
            level = self.levels[levelNo]
            for client in level.clients:
                clientLevelPairs.append((client, (-levelNo), level))
        clientLevelPairs.sort()
        out = []
        for client, levelNo, level in clientLevelPairs:
            for subRule in client.iptables_match_options():
                matchArgs = []
                for matchOption in subRule: 
                    matchArgs += matchOption.split()
                out.append(['iptables', '-A', 'FORWARD'] + matchArgs + ['-j', level.chain_name()])
        return out

    def iptables_script(self):
        blocks = [self.iptables_authorize_clients()]

        out = ""
        for level in self.levels:
            blocks.append(level.iptables_server_chain())

        for block in blocks:
            for command in block:
                out += " ".join(command) + "\n"
            out += "\n"
        return out


if __name__ == '__main__':
    import sys
    policyFile = open(sys.argv[1])
    policy = Policy.from_yaml(policyFile)
    print policy.iptables_script()

